Skip to content

fix(docker): install torchvision in GPU load-balancer image - #113

Merged
deanq merged 1 commit into
mainfrom
fix/lb-image-torchvision
Oct 27, 2026
Merged

deanq merged 1 commit into
mainfrom
fix/lb-image-torchvision

Conversation

@justinwlin

@justinwlin justinwlin commented Oct 15, 2026 •

Copy link
Copy Markdown
Collaborator

Installs the pinned torchvision in Dockerfile-lb, matching what #101 did for Dockerfile.

Why

flash build never bundles torchvision, because a bundled copy would shadow one matched to the image's torch. The queue-based GPU image has installed torchvision==0.24.1 since #101. The load-balanced GPU image never got it, so LB endpoints that import torchvision fail with ModuleNotFoundError.

Change

  • ARG TORCHVISION_VERSION=0.24.1, the release paired with the base torch 2.9.1.
  • Install from the cu128 index, the same install block as Dockerfile.
  • Add import torchvision to the build-time verify step.
Verification
  • The install block is identical to Dockerfile (diff of the two blocks is empty).
  • The published images confirm the gap: runpod/flash:py3.12-latest has the torchvision step in its build history, and runpod/flash-lb:py3.12-latest doesn't.
  • torch==2.9.1+cu128 + torchvision==0.24.1 resolves cleanly against the cu128 index. PyPI's latest torchvision 0.29.1 needs torch>=2.14, which is why the pin matters.
  • I didn't run a local image build (amd64 GPU image). CI builds Dockerfile-lb for Python 3.10 through 3.13, and the non-3.12 path reinstalls torch before this step.

Related Flash PRs

runpod-flash stops getting fixes on 2026-10-29, so the runpod/flash changes need to merge and release before then.

PR What it does Status (2026-10-27) Requires before merge
#113 Installs torchvision in the GPU load-balancer image Merged, not released. Ships with release PR #112 (1.6.1). Done.
runpod/flash#390 Keeps torch's dependency tree out of flash build archives Approved #112 merged so the 1.6.1 worker image is released.
runpod/flash#394 Waits for cold-start jobs instead of failing at 90s, and prints a working curl example after flash deploy Approved, review round 1 addressed Nothing.
runpod/docs#923 Matches the Flash quickstart and tutorials to a live run on 1.21.0 Approved, review round 1 addressed Nothing. Documents 1.21.0 as released. Revisit the 90-second guidance once runpod/flash#394 ships.
runpod/flash#369 Reads runpodctl's top-level apikey as a credentials fallback Draft Nothing.

Dockerfile got a pinned torchvision in #101 (SLS-377), but Dockerfile-lb did
not. flash build auto-excludes torchvision from tarballs, so load-balanced GPU
endpoints that import torchvision fail with ModuleNotFoundError.

Mirror the Dockerfile install: torchvision 0.24.1 from the cu128 index, paired
with the base image torch 2.9.1, plus the build-time import check.
@deanq
deanq merged commit 95de97c into main Oct 27, 2026
25 checks passed
@deanq
deanq deleted the fix/lb-image-torchvision branch October 27, 2026 17:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants